#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
从线上生产数据库导出训练数据
功能：从 mars_training_mode_chat 表提取 ID > 14 的数据并保存为 JSON
"""

import json
from sqlalchemy import create_engine, Column, Integer, String, TEXT, DateTime
from sqlalchemy.orm import sessionmaker, declarative_base
from loguru import logger

# 数据库配置（线上生产环境）
DB_CONFIG = {
    "host": "47.113.120.248",
    "port": 3306,
    "user": "root",
    "password": "Mars_nginx_chat_llm",
    "database": "mars_message_info_prod",
    "charset": "utf8mb4"
}

# 创建数据库连接
DATABASE_URL = f"mysql+pymysql://{DB_CONFIG['user']}:{DB_CONFIG['password']}@{DB_CONFIG['host']}:{DB_CONFIG['port']}/{DB_CONFIG['database']}?charset={DB_CONFIG['charset']}"

Base = declarative_base()

class MarsTrainModeKBChatMessageModel(Base):
    """训练模式聊天记录模型"""
    __tablename__ = 'mars_training_mode_chat'
    
    id = Column(Integer, primary_key=True, autoincrement=True)
    user_id = Column(String(64))
    pet_query = Column(TEXT)
    query = Column(TEXT)
    conversation_id = Column(String(128))
    chat_type = Column(String(50))
    instruction = Column(String(50))
    knowledge_base_name = Column(String(50))
    create_time = Column(DateTime(timezone=True))


def export_training_data(output_file="datas.json", min_id=14):
    """
    导出训练数据
    
    Args:
        output_file: 输出文件路径
        min_id: 最小ID，只导出大于此ID的数据
    """
    try:
        # 创建数据库引擎
        logger.info(f"正在连接数据库 {DB_CONFIG['host']}:{DB_CONFIG['port']}/{DB_CONFIG['database']}...")
        engine = create_engine(DATABASE_URL, echo=False)
        
        # 创建会话
        Session = sessionmaker(bind=engine)
        session = Session()
        
        # 查询数据 (ID > min_id)
        logger.info(f"正在查询 mars_training_mode_chat 表 (id > {min_id})...")
        records = session.query(MarsTrainModeKBChatMessageModel).filter(
            MarsTrainModeKBChatMessageModel.id > min_id
        ).order_by(MarsTrainModeKBChatMessageModel.id).all()
        
        if not records:
            logger.warning(f"未找到 ID > {min_id} 的数据")
            session.close()
            return
        
        logger.info(f"成功查询到 {len(records)} 条记录")
        
        # 转换为JSON格式
        training_data = []
        for record in records:
            data_item = {
                "user_id": record.user_id or "",
                "pet_query": record.pet_query or "",
                "query": record.query or ""
            }
            training_data.append(data_item)
            logger.debug(f"ID: {record.id} | user_id: {record.user_id} | pet_query长度: {len(record.pet_query or '')} | query长度: {len(record.query or '')}")
        
        # 保存到JSON文件
        logger.info(f"正在保存到文件: {output_file}...")
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(training_data, f, ensure_ascii=False, indent=4)
        
        logger.success(f"✅ 成功导出 {len(training_data)} 条训练数据到 {output_file}")
        
        # 显示统计信息
        logger.info("=" * 60)
        logger.info(f"导出统计:")
        logger.info(f"  总记录数: {len(training_data)}")
        logger.info(f"  起始ID: {records[0].id}")
        logger.info(f"  结束ID: {records[-1].id}")
        
        # 统计用户数量
        unique_users = set(r.user_id for r in records if r.user_id)
        logger.info(f"  涉及用户数: {len(unique_users)}")
        
        # 统计有pet_query的记录
        with_pet_query = sum(1 for r in records if r.pet_query and r.pet_query.strip())
        logger.info(f"  有宠物提问: {with_pet_query} 条")
        logger.info(f"  无宠物提问: {len(records) - with_pet_query} 条")
        logger.info("=" * 60)
        
        # 关闭会话
        session.close()
        
    except Exception as e:
        logger.error(f"❌ 导出失败: {type(e).__name__} - {str(e)}")
        raise


if __name__ == "__main__":
    # 导出 ID > 14 的数据到 datas.json
    export_training_data(
        output_file="datas.json",
        min_id=14
    )
    
    logger.info("✨ 导出完成！可以运行 json_trainer.py 进行训练了")

